import jax
import jax.numpy as jnp
import equinox as eqx
from jaxtyping import Array, Float, Int, Bool, PyTree
from .modules.utils import forward_pass, get_dts # , pad_to_len
from .modules.transformer import Transformer
from .modules.temporal_enc import ShiftedTemporalEnc
from .modules.ode import integrate

class SAHPTransformer(eqx.Module):
    transformer: Transformer
    temporal_enc: ShiftedTemporalEnc

    def __init__(self, hdim: int, embed_dim: int, nlayers: int, nhead: int, dropout_rate: float, key: Array):
        key1, key2 = jax.random.split(key)
        self.transformer = Transformer(embed_dim, nlayers, nhead, hdim, hdim, 2*hdim, dropout_rate, key1)
        self.temporal_enc = ShiftedTemporalEnc(embed_dim, 2000, key2)

    def __call__(self, dts: Float[Array, "T"], embeds: Float[Array, "T embed_dim"], key: Array):
        ts = jnp.cumsum(dts) # (T, )
        t_enc = self.temporal_enc(ts) # (T, embed_dim)
        hs = self.transformer(embeds, t_enc, key) # (T, hdim)
        return hs

class SAHP(eqx.Module):
    enc: SAHPTransformer
    mu_eta_gamma: eqx.Module
    intensity_net: list
    embed: eqx.Module

    def __init__(self, hdim: int, num_types: int, embed_dim: int, nlayers: int, nhead: int, key: Array):
        key1, key2, key3, key4 = jax.random.split(key, 4)
        self.enc = SAHPTransformer(hdim, embed_dim, nlayers, nhead, .1, key1)
        self.mu_eta_gamma = eqx.nn.Linear(hdim, 3*hdim, key=key2)
        self.intensity_net = [
            eqx.nn.Linear(hdim, num_types, key=key3),
            jax.nn.softplus
        ]
        self.embed = eqx.nn.Embedding(num_types, embed_dim, key=key4)

    def _lls(self, ts_arg, marks_arg, mask_arg, key_arg):
        dts = get_dts(ts_arg)
        dts = jnp.where(mask_arg, dts, 0.)
        embeds = jax.vmap(self.embed)(marks_arg)
        hs = self.enc(dts, embeds, key_arg) # (max_len, hdim)
        ll, dt_ll, mark_ll = jax.vmap(self._get_ll)(hs[:-1], dts[1:], marks_arg[1:], mask_arg[1:])
        ll = ll.sum()
        dt_ll = dt_ll.sum()
        mark_ll = mark_ll.sum()
        return ll, dt_ll, mark_ll

    @eqx.filter_jit
    def __call__(self, ts: Float[Array, "T"], marks: Int[Array, "T"], mask: Bool[Array, "T"], key=None):
        ll, dt_ll, mark_ll = self._lls(ts, marks, mask, key)
        return ll, dt_ll, mark_ll

    def _get_intensities(self, h: Float[Array, "hdim"], dt: Float[Array, ""]):
        hdim = h.shape[0]
        param = self.mu_eta_gamma(h) # (3*hdim, )
        mu_eta = jax.nn.gelu(param[:-hdim]) # (2*hdim, )
        gamma = jax.nn.softplus(param[-hdim:]) # (hdim, )
        mu, eta = jnp.split(mu_eta, 2)
        temp = jax.nn.tanh(mu + (eta-mu) * jnp.exp(-gamma*dt))
        intensities = forward_pass(self.intensity_net, temp)
        return intensities

    def _get_ll(self, h: Float[Array, "hdim"], dt: Float[Array, ""], mark: Int[Array, ""], mask: Bool[Array, ""]):
        sample_ratios = jnp.linspace(0., 1., 100) # (100, )
        sampled_dts = sample_ratios * dt  # (100, )
        intensities = jax.vmap(self._get_intensities, (None, 0))(h, sampled_dts) # (100, num_types)
        non_intensity = intensities[:-1].sum(-1) # (99, )
        Lambda = non_intensity.mean() * dt # ()
        event_intensities = intensities[-1] # (num_types)
        event_intensity = event_intensities[mark] # ()
        ll = jnp.log(event_intensity+1e-6) - Lambda
        mark_dist = event_intensity / (event_intensities.sum()+1e-6)
        mark_ll = jnp.log(mark_dist+1e-6)
        dt_ll = ll - mark_ll
        ll = jnp.where(mask, ll, 0.)
        dt_ll = jnp.where(mask, dt_ll, 0.)
        mark_ll = jnp.where(mask, mark_ll, 0.)
        return ll, dt_ll, mark_ll

    def intensity_at(self, h: Float[Array, "hdim"], dt: Float[Array, ""]):
        intensities = self._get_intensities(h, dt)
        intensity = intensities.sum()
        return intensity

    def _predict(self, h: Float[Array, "hdim"], dt_max: Float[Array, ""]):
        Lambda0 = jnp.asarray(0., dtype=jnp.float32)
        int_tf0 = jnp.asarray(0., dtype=jnp.float32)
        H0 = (Lambda0, int_tf0)
        def func(t, H, args=None):
            Lambda, int_tf = H
            intensity = self.intensity_at(h, t)
            f_dt = intensity * jnp.exp(-Lambda)
            tf = t * f_dt
            return intensity, tf
        _, Efdt = integrate(func, 0., dt_max, H0, None)
        intensities = self._get_intensities(h, Efdt) # (num_types, )
        mark_predict = intensities.argmax(-1)
        return Efdt, mark_predict
    
    @eqx.filter_jit
    def rolling_predict(self, ts: Float[Array, "T"], marks: Int[Array, "T"], mask: Bool[Array, "T"], dt_max: float, key=None):
        dts = get_dts(ts)
        dts = jnp.where(mask, dts, 0.)
        embeds = jax.vmap(self.embed)(marks)
        hs = self.enc(dts, embeds, key) # (T, hdim)
        dt_predict, mark_predict = jax.vmap(self._predict, (0, None))(hs[:-1], dt_max)
        return (dt_predict, mark_predict), (dts[1:], marks[1:]), mask[1:]
    
    @eqx.filter_jit
    def encode(self, ts: Float[Array, "T"], marks: Int[Array, "T"], mask: Bool[Array, "T"], key=None):
        dts = get_dts(ts)
        dts = jnp.where(mask, dts, 0.)
        embeds = jax.vmap(self.embed)(marks)
        hs = self.enc(dts, embeds, key) # (T, hdim)
        return hs

    @eqx.filter_jit
    def intensities_at(self, h: Float[Array, "hdim"], dts: Float[Array, "num_samples"], key=None):
        intensities = jax.vmap(self._get_intensities, (None, 0))(h, dts) # (num_samples, num_types)
        return intensities